# -*- coding:utf-8 -*-

# @Time    : 2023/4/20 02:09
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : construct_graph.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
import pandas as pd
class TrieNode:
    def __init__(self, name):
        self.name = name
        self.children = {}
        self.is_end_of_word = False

class Trie:
    def __init__(self):
        self.root = TrieNode(None)

    def insert(self, tag_list):
        node = self.root
        for tag in tag_list:
            if tag not in node.children:
                node.children[tag] = TrieNode(tag)
            node = node.children[tag]
        node.is_end_of_word = True

    def search(self, tag_list):
        node = self.root
        for tag in tag_list:
            if tag not in node.children:
                return False
            node = node.children[tag]
        return node.is_end_of_word

    def starts_with(self, prefix):
        node = self.root
        for tag in prefix:
            if tag not in node.children:
                return False
            node = node.children[tag]
        return True

    # 打印tire树
    def print_trie(self):
        self.print_trie_node(self.root, 0)

    def print_trie_node(self, node, level):
        if node.is_end_of_word:
            print(" " * level + node.name + " is end of word")
        else:
            if node.name:
                print(" " * level + node.name)
            else:
                print(" " * level + "root")
        for child in node.children.values():
            self.print_trie_node(child, level + 1)

    # 通过pygraphviz绘制tire树
    def draw_trie(self):
        import pygraphviz as pgv
        G = pgv.AGraph(directed=True, strict=True)
        max_depth = 3
        self.draw_trie_node(self.root, G, 0 , max_depth)
        G.layout(prog='dot')
        G.draw('png/trie.png')

    def draw_trie_node(self, node, G, depth_index, max_depth):
        if depth_index > max_depth:
            return
        if node.name:
            G.add_node(node.name)
        for child in node.children.values():
            G.add_edge(node.name, child.name)
            self.draw_trie_node(child, G, depth_index + 1, max_depth)


# 构建tire树类
class ConstructGraph(object):
    def __init__(self):
        self.trie = Trie()

    def convert_csv_to_dialogue_list(self, path):
        datas = pd.read_csv(path)
        dialogue_list = []
        dialogue = []
        session_id = ""
        for index, row in datas.iterrows():
            if (row["session_id"] != session_id):
                if len(dialogue) > 0:
                    dialogue_list.append(dialogue)
                    dialogue = []
                session_id = row["session_id"]
                dialogue.append(dict(row))
            else:
                dialogue.append(dict(row))
        return dialogue_list

    def construct_graph(self):
        datas = self.convert_csv_to_dialogue_list("data/tag_dialogue_data_with_strategy.csv")
        for dialogue in datas:
            prev_role = -1 # -1:用户 1:机器人
            node_list = []
            for row in dialogue:
                if row["predict_tag"] and pd.notna(row["predict_tag"]):
                    if row["role"] == "机器人":
                        now_role = 1
                    else:
                        now_role = -1

                    if prev_role != now_role:
                        node_list.append(row["predict_tag"])
                        prev_role = now_role
                    else:
                        break
            if len(node_list) > 0:
                self.trie.insert(node_list)

        self.trie.print_trie()
        # self.trie.draw_trie()

if __name__ == '__main__':
    construct_graph = ConstructGraph()
    construct_graph.construct_graph()




